/*
 * Copyright 2010 Nicholas Alexander Cross
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *       http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.dataporter.io.jdbc;

import org.apache.log4j.Logger;
import org.dataporter.io.DataEntity;
import org.dataporter.io.DataModel;

import java.sql.Connection;
import java.sql.DatabaseMetaData;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.*;

public class JDBCModel implements DataModel {

    private static final Logger logger = Logger.getLogger( JDBCEntity.class );

    private static final String TABLE = "TABLE";
    private static final String TABLE_NAME = "TABLE_NAME";
    private static final String TABLES_PATTERN = "%";

    private static final String COLUMN_NAME = "COLUMN_NAME";
    private static final String COLUMNS_PATTERN = "%";

    private final Map<String, JDBCEntity> entities = new HashMap<String, JDBCEntity>();

    public JDBCModel( Connection connection, String schema ) throws SQLException {
        this( connection, schema, null );
    }

    public JDBCModel( Connection connection, String schema, Set<String> allowedTables ) throws SQLException {
        String[] types = {TABLE};

        DatabaseMetaData metaData = connection.getMetaData();
        ResultSet tables = metaData.getTables( null, schema, TABLES_PATTERN, types );

        while ( tables.next() ) {
            String tableName = tables.getString( TABLE_NAME );

            if ( allowedTables == null || allowedTables.contains( tableName ) ) {
                List<String> columnNames = new ArrayList<String>();
                List<String> primaryKeyNames = new ArrayList<String>();

                ResultSet columns = metaData.getColumns( null, schema, tableName, COLUMNS_PATTERN );
                ResultSet primaryKeys = metaData.getPrimaryKeys( null, schema, tableName );

                while ( columns.next() ) {
                    String columnName = columns.getString( COLUMN_NAME );
                    columnNames.add( columnName );
                }

                while ( primaryKeys.next() ) {
                    String primaryKeyName = primaryKeys.getString( COLUMN_NAME );
                    primaryKeyNames.add( primaryKeyName );
                }

                entities.put( tableName, new JDBCEntity( tableName, primaryKeyNames, columnNames ) );
            }
        }
    }

    public DataEntity getEntity( String entityName ) {
        return entities.get( entityName );
    }

    public Set<String> entityNames() {
        return entities.keySet();
    }

    private static class JDBCEntity implements DataEntity {

        private final String entityName;
        private final List<String> primaryKeyNames;
        private final List<String> columnNames;

        public JDBCEntity( String entityName, List<String> primaryKeyNames, List<String> columnNames ) {
            this.entityName = entityName;
            this.columnNames = columnNames;
            this.primaryKeyNames = primaryKeyNames;
        }

        public String entityName() {
            return entityName;
        }

        public String batchFetchKeyName() {
            return primaryKeyNames.get( 0 );
        }

        public List<String> getPrimaryKeyNames() {
            return primaryKeyNames;
        }

        public List<String> getColumnNames() {
            return columnNames;
        }
    }
}
